/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package pl.us.ij.plugins.segmentation.commands.watershed;

import ij.process.ImageProcessor;

import java.awt.*;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;

/**
 *
 * @author Przemek
 */
public class WatershedGrayscale {

    private WatershedPixel FictitiousPixel = new WatershedPixel();
    private int _currentLabel = 0;
    private int _currentDistance = 0;
    private FifoQueue _fifoQueue = new FifoQueue();
    // each pixel can be accesesd from 2 places: a dictionary for faster direct lookup of neighbouring pixels
    // or from a height ordered list
    // sorted array of pixels according to height
    private List<List<WatershedPixel>> _heightSortedList;
    // String in the form "X,Y" is used as a key for the dictionary lookup of a pixel
    private HashMap<String, WatershedPixel> _pixelMap;
    private int _watershedPixelCount = 0;
    private int _numberOfNeighbours = 8;
    private boolean _borderInWhite;
    private int _pictureWidth = 0;
    private int _pictureHeight = 0;

    public WatershedGrayscale() throws Exception {
        this(8);
    }

    public WatershedGrayscale(int numberOfNeighbours) throws Exception {
        if (numberOfNeighbours != 8 && numberOfNeighbours != 4) {
            throw new Exception("Invalid number of neighbour pixels to check. Valid values are 4 and 8.");
        }
        _borderInWhite = true;
        _numberOfNeighbours = numberOfNeighbours;
        _heightSortedList = new LinkedList<List<WatershedPixel>>();//256

        for (int i = 0; i < 256; i++) {
            _heightSortedList.add(new LinkedList<WatershedPixel>());
        }
    }

    public void processFilter(ImageProcessor ip) {
        createPixelMapAndHeightSortedArray(ip);
        segment();
        drawWatershedLines(ip);
    }

    public int getNumberOfNeighbours() {
        return _numberOfNeighbours;
    }

    public void setNumberOfNeighbours(int value) throws Exception {

        if (value != 8 && value != 4) {
            throw new Exception("Invalid number of neighbour pixels to check. Valid values are 4 and 8.");
        }
        _numberOfNeighbours = value;
    }

    public int getLabelCount() {
        return _currentLabel;
    }

    public void setLabelCount(int value) {
        _currentLabel = value;
    }

    public boolean getBorderInWhite() {
        return _borderInWhite;
    }

    public void setBorderInWhite(boolean value) {
        _borderInWhite = value;
    }

    private void createPixelMapAndHeightSortedArray(ImageProcessor ip) {
        _pictureWidth = ip.getRoi().width;
        _pictureHeight = ip.getRoi().height;

        // pixel map holds every pixel thus size of (_pictureWidth * _pictureHeight)
        _pixelMap = new HashMap<String, WatershedPixel>(_pictureWidth * _pictureHeight);

        Rectangle r = ip.getRoi();

        // get histogram of all values in grey = height
        for (Integer y = 0; y < r.height; y++) {
            for (Integer x = 0; x < r.width; x++) {
                WatershedPixel p = new WatershedPixel(x, y, ip.get(x, y));
                // add every pixel to the pixel map
                _pixelMap.put(new Integer(p.X).toString() + "," + new Integer(p.Y).toString(), p);
                _heightSortedList.get(ip.get(x, y)).add(p);
            }
        }

        this._currentLabel = 0;
    }

    private void segment() {
        // Geodesic SKIZ (skeleton by influence zones) of each level height
        for (int h = 0; h < _heightSortedList.size(); h++) {
            // get all pixels for current height
            for (int i = 0; i < _heightSortedList.get(h).size(); i++) {
                WatershedPixel heightSortedPixel = _heightSortedList.get(h).get(i);

                heightSortedPixel.Label = WatershedCommon.MASK;
                // for each pixel on current height get neighbouring pixels
                List<WatershedPixel> neighbouringPixels = GetNeighbouringPixels(heightSortedPixel);
                // initialize queue with neighbours at level h of current basins or watersheds
                for (int j = 0; j < neighbouringPixels.size(); j++) {
                    WatershedPixel neighbouringPixel = neighbouringPixels.get(j);

                    if (neighbouringPixel.Label > 0 || neighbouringPixel.Label == WatershedCommon.WSHED) {
                        heightSortedPixel.Distance = 1;
                        _fifoQueue.AddToEnd(heightSortedPixel);
                        break;
                    }
                }
            }
            _currentDistance = 1;
            _fifoQueue.AddToEnd(FictitiousPixel);
            // extend basins
            while (true) {
                WatershedPixel p = _fifoQueue.RemoveAtFront();
                if (p.isEquals(FictitiousPixel)) {
                    if (_fifoQueue.isEmpty()) {
                        break;
                    } else {
                        _fifoQueue.AddToEnd(FictitiousPixel);
                        _currentDistance++;
                        p = _fifoQueue.RemoveAtFront();
                    }
                }
                List<WatershedPixel> neighbouringPixels = GetNeighbouringPixels(p);
                // labelling p by inspecting neighbours
                for (int j = 0; j < neighbouringPixels.size(); j++) {
                    WatershedPixel neighbouringPixel = neighbouringPixels.get(j);
                    // neighbouringPixel belongs to an existing basin or to watersheds
                    // in the original algorithm the condition is:
                    //   if (neighbouringPixel.Distance < currentDistance &&
                    //      (neighbouringPixel.Label > 0 || neighbouringPixel.Label == WatershedCommon.WSHED))
                    //   but this returns incomplete borders so the this one is used
                    if (neighbouringPixel.Distance <= _currentDistance &&
                            (neighbouringPixel.Label > 0 || neighbouringPixel.Label == WatershedCommon.WSHED)) {
                        if (neighbouringPixel.Label > 0) {
                            // the commented condition is also in the original algorithm
                            // but it also gives incomplete borders
                            if (p.Label == WatershedCommon.MASK /*|| p.Label == WatershedCommon.WSHED*/) {
                                p.Label = neighbouringPixel.Label;
                            } else if (p.Label != neighbouringPixel.Label) {
                                p.Label = WatershedCommon.WSHED;
                                _watershedPixelCount++;
                            }
                        } else if (p.Label == WatershedCommon.MASK) {
                            p.Label = WatershedCommon.WSHED;
                            _watershedPixelCount++;
                        }
                    } // neighbouringPixel is plateau pixel
                    else if (neighbouringPixel.Label == WatershedCommon.MASK && neighbouringPixel.Distance == 0) {
                        neighbouringPixel.Distance = _currentDistance + 1;
                        _fifoQueue.AddToEnd(neighbouringPixel);
                    }
                }
            }
            // detect and processImage new minima at height level h
            for (int i = 0; i < _heightSortedList.get(h).size(); i++) {
                WatershedPixel p = _heightSortedList.get(h).get(i);

                // reset distance to zero
                p.Distance = 0;
                // if true then p is inside a new minimum
                if (p.Label == WatershedCommon.MASK) {
                    // create new label
                    _currentLabel++;
                    p.Label = _currentLabel;
                    _fifoQueue.AddToEnd(p);
                    while (!_fifoQueue.isEmpty()) {
                        WatershedPixel q = _fifoQueue.RemoveAtFront();
                        // check neighbours of q
                        List<WatershedPixel> neighbouringPixels = GetNeighbouringPixels(q);
                        for (int j = 0; j < neighbouringPixels.size(); j++) {
                            WatershedPixel neighbouringPixel = neighbouringPixels.get(j);

                            if (neighbouringPixel.Label == WatershedCommon.MASK) {
                                neighbouringPixel.Label = _currentLabel;
                                _fifoQueue.AddToEnd(neighbouringPixel);
                            }
                        }
                    }
                }
            }
        }
    }

    private List<WatershedPixel> GetNeighbouringPixels(WatershedPixel centerPixel) {
        List<WatershedPixel> temp = new LinkedList<WatershedPixel>();
        if (_numberOfNeighbours == 8) {
            /*
            CP = Center pixel
            (X,Y) -- get all 8 connected
            |-1,-1|0,-1|1,-1|
            |-1, 0| CP |1, 0|
            |-1,+1|0,+1|1,+1|
             */
            // -1, -1
            if ((centerPixel.X - 1) >= 0 && (centerPixel.Y - 1) >= 0) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X - 1)).toString() + "," + new Integer((centerPixel.Y - 1)).toString()));
            }
            //  0, -1
            if ((centerPixel.Y - 1) >= 0) {
                temp.add(_pixelMap.get(new Integer(centerPixel.X).toString() + "," + new Integer((centerPixel.Y - 1)).toString()));
            }
            //  1, -1
            if ((centerPixel.X + 1) < _pictureWidth && (centerPixel.Y - 1) >= 0) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X + 1)).toString() + "," + new Integer((centerPixel.Y - 1)).toString()));
            }
            // -1, 0
            if ((centerPixel.X - 1) >= 0) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X - 1)).toString() + "," + new Integer(centerPixel.Y).toString()));
            }
            //  1, 0
            if ((centerPixel.X + 1) < _pictureWidth) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X + 1)).toString() + "," + new Integer(centerPixel.Y).toString()));
            }
            // -1, 1
            if ((centerPixel.X - 1) >= 0 && (centerPixel.Y + 1) < _pictureHeight) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X - 1)).toString() + "," + new Integer((centerPixel.Y + 1)).toString()));
            }
            //  0, 1
            if ((centerPixel.Y + 1) < _pictureHeight) {
                temp.add(_pixelMap.get(new Integer(centerPixel.X).toString() + "," + new Integer((centerPixel.Y + 1)).toString()));
            }
            //  1, 1
            if ((centerPixel.X + 1) < _pictureWidth && (centerPixel.Y + 1) < _pictureHeight) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X + 1)).toString() + "," + new Integer((centerPixel.Y + 1)).toString()));
            }
        } else {
            /*
            CP = Center pixel, N/A = not used
            (X,Y) -- get only 4 connected
            | N/A |0,-1| N/A |
            |-1, 0| CP |+1, 0|
            | N/A |0,+1| N/A |
             */
            //  -1, 0
            if ((centerPixel.X - 1) >= 0) {
                temp.add(_pixelMap.get(new Integer((centerPixel.X - 1)).toString() + "," + new Integer(centerPixel.Y).toString()));
            }
            //  0, -1
            if ((centerPixel.Y - 1) >= 0) {
                temp.add(_pixelMap.get(new Integer(centerPixel.X).toString() + "," + new Integer(centerPixel.Y - 1).toString()));
            }
            //  1, 0
            if ((centerPixel.X + 1) < _pictureWidth) {
                temp.add(_pixelMap.get(new Integer(centerPixel.X + 1).toString() + "," + new Integer(centerPixel.Y).toString()));
            }
            //  0, 1
            if ((centerPixel.Y + 1) < _pictureHeight) {
                temp.add(_pixelMap.get(new Integer(centerPixel.X).toString() + "," + new Integer(centerPixel.Y + 1).toString()));
            }
        }
        return temp;
    }

    private void drawWatershedLines(ImageProcessor ip) {
        if (_watershedPixelCount == 0) {
            return;
        }

        int watershedColor = 255;
        if (!_borderInWhite) {
            watershedColor = 0;
        }

        Rectangle r = ip.getRoi();

        for (Integer y = 0; y < r.height; y++) {
            for (Integer x = 0; x < r.width; x++) {
                // if the pixel in our map is watershed pixel then draw it
                if (_pixelMap.get(x.toString() + "," + y.toString()).Label == WatershedCommon.WSHED) {
                    ip.set(x, y, watershedColor);
                }
            }
        }
    }
}

   


